package com.qzdatasoft.ratelimit.core;

import com.qzdatasoft.HierarchicalRateLimiter;
import com.qzdatasoft.ratelimit.adapter.ClientInfoExtractor;
import com.qzdatasoft.ratelimit.adapter.RateLimitRuleCacheManager;
import jakarta.servlet.FilterChain;
import jakarta.servlet.ServletException;
import jakarta.servlet.http.HttpServletRequest;
import jakarta.servlet.http.HttpServletResponse;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.web.filter.OncePerRequestFilter;
import com.qzdatasoft.autoconfigure.ratelimit.RateLimitProperties;

import java.io.IOException;

public class RateLimitFilter extends OncePerRequestFilter {
    private static final Logger logger = LoggerFactory.getLogger(RateLimitFilter.class);
    private final HierarchicalRateLimiter limiter;
//    private final RateLimitProperties props;

    private final ClientInfoExtractor clientInfoExtractor;

    public RateLimitFilter(HierarchicalRateLimiter limiter,
                           //RateLimitProperties props,
                           ClientInfoExtractor clientInfoExtractor
                           ) {
        this.limiter = limiter;
//        this.props = props;

        this.clientInfoExtractor = clientInfoExtractor;

        logger.info("{} initialized.", this.getClass().getSimpleName());
    }

    @Override
    protected void doFilterInternal(HttpServletRequest request, HttpServletResponse response, FilterChain filterChain)
            throws ServletException, IOException {
        String ip = clientInfoExtractor.extractClientIp(request);
        String user = clientInfoExtractor.extractUserId(request);
        String url = request.getRequestURI();

        RequestTask task = new RequestTask(request, response, filterChain);
        //task.setLimiter(limiter);
        boolean enqueued = limiter.tryAcquire(task, url, ip, user);
        if (!enqueued) {
            response.setStatus(429);
            response.getWriter().write("Too Many Requests");
            return;
        }

        boolean allowed = false;
        try {
            allowed = task.awaitAllowed(/*props.getTimeoutMillis()*/);
        } catch (InterruptedException e) {
            Thread.currentThread().interrupt();
        }

        if (!allowed) {
            // try to remove in case still in queue
            LeakyBucketLimiter limiter = task.getLimiter();
            if (limiter != null) {
                limiter.removeTask(task);
            }
            //task.getLimiter().removeTask(task);
            response.setStatus(429);
            response.getWriter().write("Too Many Requests (timeout)");
            return;
        }

        // 被放行，当前线程继续执行下游链
        filterChain.doFilter(request, response);
    }
}

